Skip to content

Fourier-domain BANE#206

Open
AlecThomson wants to merge 40 commits into
PaulHancock:mainfrom
AlecThomson:dask
Open

Fourier-domain BANE#206
AlecThomson wants to merge 40 commits into
PaulHancock:mainfrom
AlecThomson:dask

Conversation

@AlecThomson
Copy link
Copy Markdown

Hi @PaulHancock,

@tjgalvin and I have been experimenting with using a Fourier approach for BANE. The motivation for this was to boost the speed, especially when running on large cubes (such as those produced for polarisation products). We were also inspired by the Fourier approach used by WSClean for background/noise estimation.

I've put together a performant script that uses Numba for some very nice speed ups. This approach can blast through large cubes in short order. I've also added some options for handling 3D and 4D cubes which might be useful for regular BANE too.

The last outstanding issue is that is seems that output images are offset up and right by the size of the convolution kernel (a.k.a box size). I'll need to do some deeper digging to figure out how to fix that. In the meantime if you have any thoughts on that, please let me know!

Anyway, I'd be really keen to get your thoughts on this kind of approach :)

@PaulHancock
Copy link
Copy Markdown
Owner

Need to include numba and radio_beam as dependencies.

@PaulHancock
Copy link
Copy Markdown
Owner

I'm getting an error when running that i'm not sure how to fix:

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Unknown attribute 'rfft2' of type Module(<module 'numpy.fft' from '/data/alpha/hancock/AlecThomson/Aegean/env/lib/python3.8/site-packages/numpy/fft/__init__.py'>)

File "AegeanTools/BANE_fft.py", line 69:
def fft_average(image: np.ndarray, kernel: np.ndarray, kern_sum: float) -> np.ndarray:
    <source elided>
    """
    image_fft = fft.rfft2(image)
    ^

During: typing of get attribute at /data/alpha/hancock/AlecThomson/Aegean/AegeanTools/BANE_fft.py (69)

File "AegeanTools/BANE_fft.py", line 69:
def fft_average(image: np.ndarray, kernel: np.ndarray, kern_sum: float) -> np.ndarray:
    <source elided>
    """
    image_fft = fft.rfft2(image)
    ^

I'm not savvy enough with numba or python Typing to know how to address this.

@AlecThomson
Copy link
Copy Markdown
Author

@PaulHancock what versions of Numpy and Numba are you rocking?

@AlecThomson
Copy link
Copy Markdown
Author

AlecThomson commented Feb 6, 2024

Ah sorry @PaulHancock, I remembered some extra steps to get this to work. Using Numba and Numpy-FFT also requires rocket-fft.

It's a little sneaky since no import is required!

So the fix is a pip install rocket-fft - I'll add this to the deps as well

@PaulHancock
Copy link
Copy Markdown
Owner

Thanks @AlecThomson it's working now.

@AlecThomson AlecThomson marked this pull request as ready for review June 13, 2025 06:12
@PaulHancock PaulHancock requested a review from Copilot June 13, 2025 06:47
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR introduces a Fourier-based implementation of BANE, adding tests, packaging updates, a CLI script, and a numba-based polynomial fit utility.

  • Added unit tests for pad_reflect, FFT smoothing, bane_fft, and robust_bane
  • Updated setup.py to include new dependencies and register the BANE_fft script
  • Created scripts/BANE_fft CLI wrapper and added AegeanTools/numba_polyfit.py

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

File Description
tests/unit/test_bane_fft.py Added fixtures and tests for pad_reflect, FFT, bane_fft, robust_bane
setup.py Bumped numpy range, added rocket-fft, numba, radio-beam, and CLI script registration
scripts/BANE_fft New entry-point script invoking the Fourier-based BANE CLI
AegeanTools/numba_polyfit.py New module providing numba-accelerated polynomial fitting
Comments suppressed due to low confidence (1)

AegeanTools/numba_polyfit.py:30

  • Variable name 'det_' is misleading—it holds the coefficients from lstsq, not a determinant. Consider renaming to 'coeffs' or 'coef'.
det_ = np.linalg.lstsq(a, b)[0]



@pytest.fixture
def ranmdom_arrays() -> RandomArrays:
Copy link

Copilot AI Jun 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fixture name 'ranmdom_arrays' appears to have a typo. Rename it to 'random_arrays' for clarity.

Suggested change
def ranmdom_arrays() -> RandomArrays:
def random_arrays() -> RandomArrays:

Copilot uses AI. Check for mistakes.


@pytest.fixture
def ranmdom_arrays_odd() -> RandomArrays:
Copy link

Copilot AI Jun 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fixture name 'ranmdom_arrays_odd' has the same typo. Rename it to 'random_arrays_odd' to match intent.

Suggested change
def ranmdom_arrays_odd() -> RandomArrays:
def random_arrays_odd() -> RandomArrays:

Copilot uses AI. Check for mistakes.
def test_bane_fft(ranmdom_arrays: RandomArrays):
gaussian_kernel_arr = gaussian_kernel(10)
for array, loc, scale in zip(
ranmdom_arrays.arrays, ranmdom_arrays.locs, ranmdom_arrays.scales, strict=False
Copy link

Copilot AI Jun 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The 'strict' parameter for zip was introduced in Python 3.10. Since the project supports >=3.8, this will error on 3.8–3.9. Consider removing 'strict=False' or raising your minimum Python to 3.10.

Suggested change
ranmdom_arrays.arrays, ranmdom_arrays.locs, ranmdom_arrays.scales, strict=False
ranmdom_arrays.arrays, ranmdom_arrays.locs, ranmdom_arrays.scales

Copilot uses AI. Check for mistakes.
Comment on lines +160 to +161
for array, loc, scale in zip(
ranmdom_arrays.arrays, ranmdom_arrays.locs, ranmdom_arrays.scales, strict=False
Copy link

Copilot AI Jun 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same issue here: 'strict=False' in zip requires Python 3.10+. Remove it or bump the Python requirement.

Suggested change
for array, loc, scale in zip(
ranmdom_arrays.arrays, ranmdom_arrays.locs, ranmdom_arrays.scales, strict=False
for array, loc, scale in zip_safe(
ranmdom_arrays.arrays, ranmdom_arrays.locs, ranmdom_arrays.scales

Copilot uses AI. Check for mistakes.
Comment on lines +10 to +11
import matplotlib.pyplot as plt

Copy link

Copilot AI Jun 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] matplotlib is only used in the demo under main. Move this import inside the main block to avoid imposing an unnecessary runtime dependency.

Suggested change
import matplotlib.pyplot as plt

Copilot uses AI. Check for mistakes.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants